import time
from enum import Enum

import numpy as np
import pygame
from pygame.joystick import JoystickType


class KeyType(Enum):
    axis = 0
    button = 1
    hat = 2
    ball = 3

class JoystickKey(Enum):
    ROLL = (KeyType.axis, 0)
    PITCH = (KeyType.axis, 1)




class Throttlekey(Enum):
    AXIS_X = (KeyType.axis, 0)
    AXIS_Y = (KeyType.axis, 1)
    R_THROTTLE = (KeyType.axis, 2)
    L_THROTTLE = (KeyType.axis, 3)
    INCR = (KeyType.axis, 4)

    AXIS_BTN = (KeyType.button, 1)
    UNNAME_BTN = (KeyType.button, 2)
    UNNAME_BTN_UP = (KeyType.button, 3)
    UNNAME_BTN_FORWARD = (KeyType.button, 4)
    UNNAME_BTN_DOWN = (KeyType.button, 5)
    UNNAME_BTN_BACKWARD = (KeyType.button, 6)
    UNNAME2_BTN_FORWARD = (KeyType.button, 7)
    UNNAME2_BTN_BACKWARD = (KeyType.button, 8)
    UNNAME3_BTN_FORWARD = (KeyType.button, 9)
    UNNAME3_BTN_BACKWARD = (KeyType.button, 10)
    UNNAME4_BTN_FORWARD = (KeyType.button, 11)
    UNNAME4_BTN_BACKWARD = (KeyType.button, 12)
    UNNAME5_BTN_FORWARD = (KeyType.button, 13)
    UNNAME5_BTN_BACKWARD = (KeyType.button, 14)
    UNNAME6_BTN = (KeyType.button, 15)
    ENG_L = (KeyType.button, 16)
    ENG_R = (KeyType.button, 17)
    MOTOR_L = (KeyType.button, 18)
    MOTOR_R = (KeyType.button, 19)
    APU_START = (KeyType.button, 20)
    SILENCE = (KeyType.button, 21)
    FLAPS_UP = (KeyType.button, 22)
    FLAPS_DOWN = (KeyType.button, 23)
    EAC = (KeyType.button, 24)
    RDR = (KeyType.button, 25)
    ENGAGE_BTN = (KeyType.button, 26)
    LASTE_PATH = (KeyType.button, 27)
    LASTE_ALT = (KeyType.button, 28)
    IGN_L = (KeyType.button, 31)
    IGN_R = (KeyType.button, 32)

    HAT = (KeyType.hat, 0)


class Rudderkey(Enum):
    KeyNone = (-1)
    rudder = (KeyType.axis, 2)


class FlightJoystick:
    def __init__(self):
        self.throttle_value = None
        self.yaw = None
        self.pitch = None
        self.psi = 0
        pygame.init()
        pygame.joystick.init()
        self.throttle_device_name = "Throttle - HOTAS Warthog"
        self.joystick_device_name = "Joystick - HOTAS Warthog"
        self.rudder_device_name = "T-Rudder"
        self.throttle_founded = False
        self.joystick_founded = False
        self.rudder_founded = False
        joystick_count = pygame.joystick.get_count()
        print("找到 {} 个手柄设备".format(joystick_count))
        # print(pygame.joystick.Joystick(0).get_name())
        if pygame.joystick.get_count() >= 1:
            self.joystick_device = None
            self.throttle = None
            self.rudder = None
            for i in range(pygame.joystick.get_count()):
                joystick = pygame.joystick.Joystick(i)
                print(joystick.get_name())
                if joystick.get_name() == self.throttle_device_name:
                    self.throttle = joystick
                    self.throttle_founded = True
                if joystick.get_name() == self.joystick_device_name:
                    self.joystick_device = joystick
                    self.joystick_founded = True
                if joystick.get_name() == self.rudder_device_name:
                    self.rudder = joystick
                    self.rudder_founded = True
            if self.joystick_device is not None:
                self.joystick_device.init()
            else:
                print("joystick not found")
            if self.throttle is not None:
                self.throttle.init()
            else:
                print("throttle not found")
            if self.rudder is not None:
                self.rudder.init()
            else:
                print("t_rudder not found")
        else:
            print("缺少必要控制器")

    def read_action(self):
        pygame.event.get()
        if self.joystick_founded:
            self.pitch = -self.read(JoystickKey.PITCH) * 0.8  # 俯仰数据
            self.yaw = self.read(JoystickKey.ROLL) * 0.1  # 滚转数据
        if self.rudder_founded:
            self.psi = self.read(Rudderkey.rudder)  # 偏航数据
        if self.throttle_founded:
            self.throttle_value = self.read(Throttlekey.L_THROTTLE)  # 油门数据
            self.throttle_value = 0.4 + (1 - 0.4) * (self.throttle_value - 1) / -2

    def read(self, key):
        joystick_instance = None
        if type(key) == JoystickKey:
            joystick_instance = self.joystick_device
        if type(key) == Throttlekey:
            joystick_instance = self.throttle
        if type(key) == Rudderkey:
            joystick_instance = self.rudder
        if joystick_instance is not None:
            return self.read_joystick(joystick_instance, key)
        else:
            print("joystick not found")
            return None

    def read_joystick(self, joystick_instance, joystick_key):
        if joystick_key.value[0] == KeyType.axis:
            return joystick_instance.get_axis(joystick_key.value[1])
        if joystick_key.value[0] == KeyType.button:
            return joystick_instance.get_button(joystick_key.value[1])
        if joystick_key.value[0] == KeyType.hat:
            return joystick_instance.get_hat(joystick_key.value[1])
        if joystick_key.value[0] == KeyType.ball:
            return joystick_instance.get_ball(joystick_key.value[1])
        print("Joystick is founded but KeyType is not found")
        return None

    def get_action(self):
        action = np.array([self.yaw, self.pitch, self.psi, self.throttle_value])
        return action

    def debuginfo(self, joystick: JoystickType):

        # print("\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n")
        print("name:", joystick.get_name(), flush=True)
        print("Axis:", joystick.get_numaxes(), flush=True)
        print("btn:", joystick.get_numbuttons(), flush=True)
        print("hat:", joystick.get_numhats(), flush=True)
        print("ball:", joystick.get_numballs(), flush=True)
        str = ""
        for i in range(joystick.get_numaxes()):
            str += joystick.get_axis(i).__str__() + " "
        str += "\n"
        for i in range(joystick.get_numbuttons()):
            str += joystick.get_button(i).__str__() + " "
        str += "\n"
        for i in range(joystick.get_numhats()):
            str += joystick.get_hat(i).__str__() + " "
        str += "\n"
        for i in range(joystick.get_numballs()):
            str += joystick.get_ball(i).__str__() + " "
        str += "\n"
        print("\r" + str, flush=True)


if __name__ == '__main__':
    j = FlightJoystick()
    while True:
        # j.read(Throttlekey.)
        # print(j.get_action())
        # j.debuginfo(j.joystick_device)
        time.sleep(0.5)
